#include "kernel_operator.h"
constexpr int32_t BUFFER_NUM = 2; 
constexpr int32_t CORE_NUM = 8; 

template<typename TYPE> class KernelEyeCustom { 
public:
    __aicore__ inline KernelEyeCustom() {}

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, const EyeCustomTilingData &tilingData)
    {
        this->tiling = tilingData;
        uint32_t blockIdx;
        this->blockIdx = AscendC::GetBlockIdx();
        matrixOrder = (blockIdx >= tiling.fullBlockNum) ? tiling.tailBlockLength : tiling.fullBlockLength;
        xGm.SetGlobalBuffer((__gm__ TYPE*)x, tiling.rowLength * tiling.columnLength );
        yGm.SetGlobalBuffer((__gm__ TYPE*)y, tiling.rowLength * tiling.columnLength );
        pipe.InitBuffer(inQueueX, BUFFER_NUM, sizeof(TYPE));
        pipe.InitBuffer(inQueueY, BUFFER_NUM, sizeof(TYPE));
    }

    __aicore__ inline void Process()
    {
        CopyIn();
        Compute();
        CopyOut();
    }

private:
    __aicore__ inline void CopyIn()
    {
        AscendC::LocalTensor<TYPE> xLocal = inQueueX.AllocTensor<TYPE>();
        AscendC::Duplicate(xLocal, static_cast<TYPE>(0), matrixOrder);
        AscendC::DumpTensor(xLocal, 0, 8);
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute()
    {
        AscendC::LocalTensor<TYPE> yLocal = inQueueY.AllocTensor<TYPE>();
        AscendC::LocalTensor<TYPE> xLocal = inQueueX.DeQue<TYPE>();
        AscendC::Adds(yLocal, xLocal, static_cast<TYPE>(1), matrixOrder);
        inQueueY.EnQue(yLocal);
        inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void CopyOut()
    {
        AscendC::LocalTensor<TYPE> yLocal = inQueueY.DeQue<TYPE>();
        if (matrixOrder == 0) {
            inQueueY.FreeTensor(yLocal);
            return;
        }
        uint32_t idx = blockIdx;
        uint64_t diagStart =
            (idx <= tiling.fullBlockNum)
                ? static_cast<uint64_t>(idx) * tiling.fullBlockLength
                : static_cast<uint64_t>(tiling.fullBlockNum) * tiling.fullBlockLength +
                  static_cast<uint64_t>(idx - tiling.fullBlockNum) * tiling.tailBlockLength;
        uint64_t globalOffset = diagStart * tiling.columnLength + 
                                (static_cast<int64_t>(diagStart) + 
                                static_cast<int64_t>(tiling.indexK));
        AscendC::DataCopyExtParams copyParams{static_cast<uint16_t>(1), static_cast<uint32_t>(sizeof(TYPE)), 0, 0, 0};
        for (uint32_t i = 0; i < matrixOrder; ++i) {
            uint64_t dst = globalOffset + static_cast<uint64_t>(i * (tiling.columnLength + 1));
            AscendC::DataCopyPad(yGm[dst], yLocal, copyParams);
        }
        inQueueY.FreeTensor(yLocal);
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> inQueueY;
    AscendC::GlobalTensor<TYPE> xGm;
    AscendC::GlobalTensor<TYPE> yGm;

    EyeCustomTilingData tiling;
    uint32_t rowLength;             // 矩阵行数
    uint32_t columnLength;          // 矩阵列数
    uint32_t diagLen;               // 对角线长度
    int32_t indexK;                 // 对角线索引
    uint32_t fullBlockLength;       // 每个核处理的对角线长度（向上取整）
    uint32_t tailBlockLength;       // 每个核处理的对角线长度（向下取整）
    uint32_t fullBlockNum;          // 处理大核的核数
    uint32_t tailBlockNum;          // 处理小核的核数
    uint32_t typeSize;              // 每个元素的字节数

    uint32_t matrixOrder;          // 当前核处理的矩阵阶数
    uint32_t blockIdx;
};

extern "C" __global__ __aicore__ void eye_custom(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tilingData, tiling);
    if(TILING_KEY_IS(0)){
        KernelEyeCustom<half> op;
        op.Init(x, y, tilingData);
        op.Process();
    }else if(TILING_KEY_IS(1)){
        KernelEyeCustom<float> op;
        op.Init(x, y, tilingData);
        op.Process();
    }else if(TILING_KEY_IS(2)){
        KernelEyeCustom<int16_t> op;
        op.Init(x, y, tilingData);
        op.Process();
    }else if(TILING_KEY_IS(3)){
        KernelEyeCustom<int32_t> op;
        op.Init(x, y, tilingData);
        op.Process();
    }else{
        return;
    }
}